import pandas as pd
import numpy as np
from sqlalchemy import create_engine

def moving_average_crossover_strategy(data, short_window=20, long_window=50):
    """移动平均线交叉策略"""
    # 计算移动平均线
    data['short_ma'] = data['close'].rolling(window=short_window).mean()
    data['long_ma'] = data['close'].rolling(window=long_window).mean()
    
    # 生成交易信号
    data['signal'] = 0
    data.loc[data['short_ma'] > data['long_ma'], 'signal'] = 1
    data.loc[data['short_ma'] < data['long_ma'], 'signal'] = 0
    
    # 计算持仓变化
    data['position'] = data['signal'].diff()
    
    return data


def price_breakout_strategy(data, window=20):
    """价格突破策略"""
    # 计算高低点
    data['high_window'] = data['high'].rolling(window=window).max()
    data['low_window'] = data['low'].rolling(window=window).min()
    
    # 生成交易信号
    data['signal'] = 0
    data.loc[data['close'] > data['high_window'].shift(1), 'signal'] = 1  # 突破高点买入
    data.loc[data['close'] < data['low_window'].shift(1), 'signal'] = 0  # 跌破低点卖出
    
    # 计算持仓变化
    data['position'] = data['signal'].diff()
    
    return data


def rsi_strategy(data, window=14, overbought=70, oversold=30):
    """RSI指标策略"""
    # 计算价格变化
    delta = data['close'].diff()
    
    # 分离上涨和下跌
    gain = delta.mask(delta < 0, 0)
    loss = -delta.mask(delta > 0, 0)
    
    # 计算平均收益上涨和下跌
    avg_gain = gain.rolling(window=window).mean()
    avg_loss = loss.rolling(window=window).mean()
    
    # 计算RSI
    rs = avg_gain / avg_loss
    data['rsi'] = 100 - (100 / (1 + rs))
    
    # 生成交易信号
    data['signal'] = 0
    data.loc[data['rsi'] < oversold, 'signal'] = 1  # RSI低于超卖线买入
    data.loc[data['rsi'] > overbought, 'signal'] = 0  # RSI高于超买线卖出
    
    # 计算持仓变化
    data['position'] = data['signal'].diff()
    
    return data


def comprehensive_strategy(data, 
                          lookback_period=10, volume_multiple=2.0,
                          short_window=5, long_window=20,
                          bollinger_window=20, num_std=2,
                          mean_reversion_window=10, z_score_threshold=1.0):
    """
    综合量化交易策略
    
    融合动量、双均线、均值回归、布林带和成交量五大因子，
    通过动态加权整合各子策略信号，突破单一策略局限：
    - 动量策略（20%）捕捉趋势惯性
    - 双均线（25%）判断趋势转折
    - 布林带（20%）识别超买超卖
    - 均值回归（15%）捕捉反转机会
    - 成交量（20%）验证价格变动有效性
    
    参数:
        data: DataFrame, 包含股票数据
        lookback_period: 成交量回溯期(天数)
        volume_multiple: 成交量倍数阈值
        short_window: 短期均线窗口
        long_window: 长期均线窗口
        bollinger_window: 布林带窗口
        num_std: 布林带标准差倍数
        mean_reversion_window: 均值回归窗口
        z_score_threshold: Z分数阈值
    
    返回:
        data: DataFrame, 添加了信号列的数据
    """
    # 复制数据以避免修改原数据
    df = data.copy()
    
    # 确保有pct_chg列，如果没有则计算
    if 'pct_chg' not in df.columns:
        df['pct_chg'] = df['close'].pct_change() * 100
    
    # 1. 动量策略
    df['momentum'] = df['pct_chg'].rolling(lookback_period).apply(
        lambda x: (x / 100 + 1).prod() - 1
    )
    df['momentum_signal'] = np.where(
        df['momentum'] > 0, 1,  # 正动量买入
        np.where(df['momentum'] < 0, -1, 0)  # 负动量卖出
    )
    
    # 2. 双均线策略
    df['short_ma'] = df['close'].rolling(window=short_window).mean()
    df['long_ma'] = df['close'].rolling(window=long_window).mean()
    df['ma_signal'] = np.where(
        df['short_ma'] > df['long_ma'], 1,  # 金叉买入
        np.where(df['short_ma'] < df['long_ma'], -1, 0)  # 死叉卖出
    )
    
    # 3. 布林带策略
    df['rolling_mean'] = df['close'].rolling(window=bollinger_window).mean()
    df['rolling_std'] = df['close'].rolling(window=bollinger_window).std()
    df['upper_band'] = df['rolling_mean'] + (df['rolling_std'] * num_std)
    df['lower_band'] = df['rolling_mean'] - (df['rolling_std'] * num_std)
    df['bollinger_signal'] = np.where(
        df['close'] < df['lower_band'], 1,  # 低于下轨买入
        np.where(df['close'] > df['upper_band'], -1, 0)  # 高于上轨卖出
    )
    
    # 4. 均值回归策略
    df['mean'] = df['close'].rolling(window=mean_reversion_window).mean()
    df['std'] = df['close'].rolling(window=mean_reversion_window).std()
    df['z_score'] = (df['close'] - df['mean']) / df['std']
    df['mean_reversion_signal'] = np.where(
        df['z_score'] < -z_score_threshold, 1,  # 低于均值买入
        np.where(df['z_score'] > z_score_threshold, -1, 0)  # 高于均值卖出
    )
    
    # 5. 成交量策略
    df['vol_ma'] = df['vol'].rolling(window=lookback_period).mean()
    df['vol_ratio'] = df['vol'] / df['vol_ma']
    df['volume_signal'] = np.where(
        (df['vol_ratio'] >= volume_multiple) & (df['pct_chg'] > 0), 1,  # 放量上涨买入
        np.where(
            (df['vol_ratio'] >= volume_multiple) & (df['pct_chg'] < 0), -1,  # 放量下跌卖出
            np.where(
                (df['vol_ratio'] < volume_multiple) & (df['pct_chg'] > 0), 0.5,  # 量平上涨持有
                0  # 其他情况观望
            )
        )
    )
    
    # 综合信号 (加权平均)
    df['combined_signal_value'] = (0.2 * df['momentum_signal'] +
                                   0.25 * df['ma_signal'] +
                                   0.2 * df['bollinger_signal'] +
                                   0.15 * df['mean_reversion_signal'] +
                                   0.2 * df['volume_signal'])
    
    # 生成交易信号
    # 规则：提示买入或强烈买入时买入，观望/卖出/强烈卖出时卖出，其他时候持仓不动
    df['signal'] = 0  # 默认不持仓
    
    # 初始信号：基于综合信号值
    df.loc[df['combined_signal_value'] > 0.5, 'signal'] = 1  # 强烈买入
    df.loc[(df['combined_signal_value'] > 0.2) & (df['combined_signal_value'] <= 0.5), 'signal'] = 1  # 买入
    df.loc[(df['combined_signal_value'] >= -0.2) & (df['combined_signal_value'] <= 0.2), 'signal'] = 0  # 持有（保持现有持仓）
    df.loc[(df['combined_signal_value'] > -0.5) & (df['combined_signal_value'] < -0.2), 'signal'] = 0  # 卖出（保持空仓）
    df.loc[df['combined_signal_value'] <= -0.5, 'signal'] = 0  # 强烈卖出（保持空仓）
    
    # 确保信号连续性：只在需要改变持仓时才产生信号变化
    # 对于已持仓的情况，只有在收到卖出信号时才平仓
    # 对于未持仓的情况，只有在收到买入信号时才开仓
    
    # 向前填充信号以保持持仓状态
    df['signal'] = df['signal'].ffill()
    
    # 计算持仓变化
    df['position'] = df['signal'].diff()
    
    # 确保第一天没有持仓变化
    if len(df) > 0:
        df.iloc[0, df.columns.get_loc('position')] = 0
    
    return df


def backtest_strategy(data, strategy_func, initial_capital=100000, **kwargs):
    """回测策略"""
    # 复制数据以避免修改原数据
    backtest_data = data.copy()
    
    # 应用策略
    backtest_data = strategy_func(backtest_data,** kwargs)
    
    # 计算每日收益率
    backtest_data['daily_return'] = backtest_data['close'].pct_change()
    
    # 计算策略收益率
    backtest_data['strategy_return'] = backtest_data['daily_return'] * backtest_data['signal'].shift(1)
    
    # 计算累计收益率
    backtest_data['cumulative_market'] = (1 + backtest_data['daily_return']).cumprod()
    backtest_data['cumulative_strategy'] = (1 + backtest_data['strategy_return']).cumprod()
    
    # 计算资金曲线 - 使用浮点数类型避免数据类型警告
    backtest_data['cash'] = initial_capital
    backtest_data['cash'] = backtest_data['cash'].astype(float)
    
    backtest_data['holdings'] = 0.0
    backtest_data['total_asset'] = initial_capital
    backtest_data['total_asset'] = backtest_data['total_asset'].astype(float)
    
    # 交易记录（含手续费后的真实盈亏）
    backtest_data['trade_pnl'] = 0.0  # 扣手续费后实际盈亏
    backtest_data['trade_return'] = 0.0  # 扣手续费后实际收益率
    backtest_data['commission'] = 0.0  # 佣金
    backtest_data['stamp_duty'] = 0.0  # 印花税
    
    # 记录完整交易的实际盈亏（扣手续费后）
    completed_trades_pnl = []
    
    # 初始化交易参数
    cash = initial_capital
    holdings = 0
    shares = 0
    last_buy_price = 0
    # 手续费参数（当前A股标准）
    commission_rate = 0.0001  # 佣金万分之一
    min_commission = 1  # 最低佣金1元
    stamp_duty_rate = 0.001  # 印花税千分之一（仅卖出时收取）
    
    for i in range(1, len(backtest_data)):
        current_price = backtest_data.iloc[i]['close']
        position_change = backtest_data.iloc[i]['position']
        
        # 更新资产价值
        if holdings > 0:
            backtest_data.iloc[i, backtest_data.columns.get_loc('holdings')] = shares * current_price
        else:
            backtest_data.iloc[i, backtest_data.columns.get_loc('holdings')] = 0
            
        backtest_data.iloc[i, backtest_data.columns.get_loc('cash')] = cash
        backtest_data.iloc[i, backtest_data.columns.get_loc('total_asset')] = cash + backtest_data.iloc[i]['holdings']
        
        # 买入逻辑（计算买入佣金）
        if position_change == 1:  # 买入信号
            if cash > 0:
                # 计算可买数量（全仓）
                buy_amount = cash
                shares = buy_amount / current_price
                # 取整数股（A股最低100股，此处简化为整数）
                shares = int(shares // 1)  # 实际应是100的倍数，此处简化
                actual_buy_amount = shares * current_price
                
                # 计算买入佣金（万分之一，最低1元）
                commission = actual_buy_amount * commission_rate
                commission = max(commission, min_commission)  # 不足1元按1元收
                
                # 扣佣金后更新现金
                cash -= (actual_buy_amount + commission)
                holdings = shares * current_price
                last_buy_price = current_price
                
                # 记录买入佣金
                backtest_data.iloc[i, backtest_data.columns.get_loc('commission')] = commission
        
        # 卖出逻辑（计算卖出佣金+印花税）
        elif position_change == -1:  # 卖出信号
            if holdings > 0:
                # 计算卖出金额
                sell_amount = shares * current_price
                
                # 计算卖出佣金（万分之一，最低1元）
                commission = sell_amount * commission_rate
                commission = max(commission, min_commission)
                
                # 计算印花税（千分之一，仅卖出时收）
                stamp_duty = sell_amount * stamp_duty_rate  # 印花税无最低，精确到分
                
                # 扣佣金和印花税后实际到账金额
                cash_received = sell_amount - commission - stamp_duty
                
                # 计算实际盈亏（扣手续费后）
                trade_pnl = cash_received - (last_buy_price * shares)
                trade_return = trade_pnl / (last_buy_price * shares) if last_buy_price > 0 else 0
                
                # 记录完整交易盈亏（扣手续费后）
                completed_trades_pnl.append(trade_pnl)
                
                # 更新资金和持仓
                cash = cash_received
                shares = 0
                holdings = 0
                
                # 记录交易数据
                backtest_data.iloc[i, backtest_data.columns.get_loc('trade_pnl')] = trade_pnl
                backtest_data.iloc[i, backtest_data.columns.get_loc('trade_return')] = trade_return
                backtest_data.iloc[i, backtest_data.columns.get_loc('commission')] = commission
                backtest_data.iloc[i, backtest_data.columns.get_loc('stamp_duty')] = stamp_duty
    
    # 计算最终资产
    final_asset = backtest_data.iloc[-1]['total_asset']
    total_profit = final_asset - initial_capital
    
    # 策略指标计算
    total_days = len(backtest_data)
    total_market_return = backtest_data['cumulative_market'].iloc[-1] - 1 if total_days > 0 else 0
    total_strategy_return = backtest_data['cumulative_strategy'].iloc[-1] - 1 if total_days > 0 else 0
    
    # 年化收益率（252个交易日）
    annual_market_return = (1 + total_market_return) **(252 / total_days) - 1 if total_days > 0 else 0
    annual_strategy_return = (1 + total_strategy_return)** (252 / total_days) - 1 if total_days > 0 else 0
    
    # 最大回撤
    def calculate_max_drawdown(series):
        if len(series) < 2:
            return 0
        rolling_max = series.cummax()
        drawdown = (series - rolling_max) / rolling_max
        return drawdown.min()
    
    max_drawdown_market = calculate_max_drawdown(backtest_data['cumulative_market'])
    max_drawdown_strategy = calculate_max_drawdown(backtest_data['cumulative_strategy'])
    
    # 胜率计算（基于扣手续费后的完整交易，过滤零盈亏）
    valid_trades = [pnl for pnl in completed_trades_pnl if pnl != 0]  # 排除零盈亏
    total_completed_trades = len(valid_trades)
    winning_trades = sum(1 for pnl in valid_trades if pnl > 0)
    win_rate = winning_trades / total_completed_trades if total_completed_trades > 0 else 0.0
    
    # 交易次数（买入+卖出，每笔完整交易计2次信号）
    trade_count = backtest_data['position'].abs().sum()
    
    # 平均每笔交易盈亏（扣手续费后）
    profitable_trades = sum(pnl for pnl in valid_trades if pnl > 0)
    losing_trades = sum(pnl for pnl in valid_trades if pnl < 0)
    avg_profit_per_trade = profitable_trades / (sum(1 for pnl in valid_trades if pnl > 0)) if (sum(1 for pnl in valid_trades if pnl > 0)) > 0 else 0
    avg_loss_per_trade = losing_trades / (sum(1 for pnl in valid_trades if pnl < 0)) if (sum(1 for pnl in valid_trades if pnl < 0)) > 0 else 0
    
    # 风险回报比
    risk_reward_ratio = abs(avg_profit_per_trade / avg_loss_per_trade) if avg_loss_per_trade != 0 else float('inf')
    
    # 结果指标
    metrics = {
        '初始资金': initial_capital,
        '最终资产': final_asset,
        '总盈利': total_profit,
        '总收益率': total_profit / initial_capital,
        '总交易日数': total_days,
        '市场总收益率': total_market_return,
        '策略总收益率': total_strategy_return,
        '市场年化收益率': annual_market_return,
        '策略年化收益率': annual_strategy_return,
        '市场最大回撤': max_drawdown_market,
        '策略最大回撤': max_drawdown_strategy,
        '胜率': win_rate,
        '交易次数': trade_count,
        '平均每笔盈利': avg_profit_per_trade,
        '平均每笔亏损': avg_loss_per_trade,
        '风险回报比': risk_reward_ratio
    }
    
    return backtest_data, metrics


def get_available_strategies():
    """获取可用的策略列表"""
    return [
        ("移动平均线交叉策略", "ma_crossover", moving_average_crossover_strategy),
        ("价格突破策略", "price_breakout", price_breakout_strategy),
        ("RSI指标策略", "rsi_strategy", rsi_strategy),
        ("综合量化策略", "comprehensive", comprehensive_strategy)
    ]


def get_strategy_by_name(strategy_name):
    """根据策略名称获取策略函数"""
    strategies_list = get_available_strategies()
    for name, code, func in strategies_list:
        if code == strategy_name:
            return name, func
    return None, None
